-
Notifications
You must be signed in to change notification settings - Fork 6.8k
[MXNET-1413] Adding Large Tensor support for sort operators #15170
Conversation
@mxnet-label-bot Add [pr-awaiting-review] |
@apeforest please review |
mxnet_op::Kernel<range_fwd, xpu>::Launch(s, batch_size * element_num, 1, 0, 1, | ||
kWriteTo, indices.dptr_); | ||
mxnet_op::Kernel<range_fwd, xpu>::Launch(s, batch_size * element_num, 1, static_cast<index_t>(0), | ||
static_cast<index_t>(1), kWriteTo, indices.dptr_); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe just using data initializer:
static_cast<index_t>(1), kWriteTo, indices.dptr_); | |
index_t{0}, index_t{1} |
c3af3fa
to
e62d8da
Compare
a9e2fff
to
cedab68
Compare
@@ -819,7 +821,11 @@ def get_large_matrix(): | |||
|
|||
# test for ret_typ=indices | |||
nd_ret_topk = mx.nd.topk(a_nd, axis=1, ret_typ="indices", k=3, is_ascend=True).asnumpy() | |||
assert nd_ret_topk.dtype == np.float32 # Test the default dtype | |||
# Test the default dtype | |||
if is_large_tensor_enabled: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
changed default type from float to int32/64 based on whether large tensor support is enabled or not
@@ -860,7 +866,10 @@ def get_large_matrix(): | |||
nd_ret_topk_val = nd_ret_topk_val.asnumpy() | |||
nd_ret_topk_ind = nd_ret_topk_ind.asnumpy() | |||
assert nd_ret_topk_val.dtype == dtype | |||
assert nd_ret_topk_ind.dtype == np.float32 | |||
if is_large_tensor_enabled: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
same as above
@apeforest please review |
Tensor<xpu, 1, index_t> workspace = | ||
ctx.requested[0].get_space_typed<xpu, 1, index_t>(Shape1(batch_size * k + batch_size), s); | ||
Tensor<xpu, 1, index_t> sel_indices = | ||
Tensor<xpu, 1, index_t>((workspace.dptr_), Shape1(batch_size * k), s); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: parenthesis around workspace.dptr_ seems redundant
CHECK(type_assign(&(*out_attrs)[1], mshadow::kInt32)) | ||
<< "Failed to set the type of ret_indices to int32."; | ||
#endif | ||
<< "Failed to set the type of ret_indices to int32."; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this message should be different when MSHADOW_USE_INT64_TENSOR_SIZE
is on.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think its better if we logged "Failed to set the type of ret_indices"
k = nd.topk(b, k=10, axis=0, dtype=np.int64) | ||
assert np.sum(k.asnumpy() == (LARGE_X - 1)) == SMALL_Y | ||
b = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) | ||
l = nd.topk(b, k=1, axis=-1, dtype=np.int64, ret_typ="value") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
please also test when ret_typ is "both" and "indices"
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point. Default is indices. I can check for "both" as well
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for making this change with iterations of updates to make the API backward compatible. LGTM overall except a few minor changes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM % comment about uninitalized variables.
@@ -605,30 +633,32 @@ void TopKBackwardImpl(const OpContext &ctx, | |||
using namespace mshadow::expr; | |||
Stream<xpu> *s = ctx.run_ctx.get_stream<xpu>(); | |||
CHECK(param.ret_typ == topk_enum::kReturnValue || param.ret_typ == topk_enum::kReturnBoth); | |||
int batch_size, element_num; // number of batches + the size of each batch | |||
size_t batch_size; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we get into the good practice of not leaving variables uninitialized?
This is verboten in many serious circumstances, see automotive, MISRA, aerospatial.... I know we are parsing the arguments later, but still.
Same thing for above.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
sounds good. I can do this for primitive types which come with default garbage values. I will refrain from intializing class objects, since they are null by default unless explicitly assigned memory.
@larroy does that sound good ?
@apeforest what are your thoughts on this ?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes I was referring to primitive types. Which class objects you referr to? When initializing a class the fields are not null, it actually depends on what's on the ctor. If there's no ctor and the type is pod is garbage.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, I agree with @larroy leaving no uninit variables is a good SE practice. Class object should have its own default constructors, if not, then the design of class need improvement.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@larroy I was talking about objects of Tensor class
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Only pods need to be initialized in any circumstance including class ctor (otherwise garbage), so answering to your original question, yes. And seems we all agree here. Objects don't need explicit = initialization as per ctor as @apeforest pointed out.
d7c78fa
to
496dd9f
Compare
tests/nightly/test_large_array.py
Outdated
b = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) | ||
l = nd.topk(b, k=1, axis=-1, dtype=np.int64, ret_typ="value") | ||
assert l.sum() == np.sum(np.arange(0, SMALL_Y)) | ||
b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why do we create b multiple times? Can we reuse one to save computation?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I can reuse 1st but 2nd still needs to be created.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
fine. don't need to do in this PR.
@mxnet-label-bot Add [pr-awaiting-merge] |
@access2rohit Can you check the CI status? If it got staled, please retrigger it again. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
Unfortunately this change has broken NightlyTestsForBinaries #15374 |
Description
ops supported: sort, topk, argmax
Checklist
Essentials
Please feel free to remove inapplicable items for your PR.